import numpy as np
import yaml
from pyquaternion import Quaternion

def quat2mat(quat):
    ''' Symbolic conversion from quaternion to rotation matrix'''
    w, x, y, z = quat
    return [[1 - 2*y*y-2*z*z, 2*x*y - 2*z*w, 2*x*z+2*y*w],
            [2*x*y+2*z*w, 1-2*x*x-2*z*z, 2*y*z-2*x*w],
            [2*x*z-2*y*w, 2*y*z+2*x*w, 1-2*x*x-2*y*y]]

def compose(T, R, Z):
    n = len(T)
    R = np.asarray(R)
    if R.shape != (n,n):
        raise ValueError('Expecting shape (%d,%d) for rotations' % (n,n))
    A = np.eye(n+1)
    ZS = np.diag(Z)
    A[:n,:n] = np.dot(R, ZS)
    A[:n,n] = T[:]
    return A

def load_yaml_file(file_path):
    with open(file_path, 'r') as f:
        ret= yaml.safe_load(f.read())
    return ret

def build_extrinsics(R, T):
    r = quat2mat(R)
    t = T
    return compose(t, r, [1.0, 1.0, 1.0])

if __name__ == '__main__':
    fp = '/home/data/1354/extrinsics/145046_fa2sd4a06W152AIR_420_1626156703_1626157310_102_obstacle.yaml'
    a = load_yaml_file(fp)
    r, t = a['transform']['rotation'], a['transform']['translation']
    e = build_extrinsics([r['w'], r['x'], r['y'], r['z']], [t['x'], t['y'], t['z']])
    print(e)
